#!/usr/bin/env python
# -*- coding: utf-8 -*-

""" By Martin Senande-Rivera
    For Spatial and temporal expansion of global wildland fire activity in response to climate change """

import os
import glob, sys
import numpy as np
import xarray as xr
import pandas as pd
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.cm import get_cmap
from mpl_toolkits.basemap import Basemap, maskoceans
import matplotlib.colors as mcolors
import datetime as dt
from matplotlib.colors import LinearSegmentedColormap
import cmaps
import cmocean
from pandas import Series, DataFrame
from mpl_toolkits.axes_grid1.inset_locator import inset_axes


font = {'size'   : 5}
matplotlib.rc('font', **font)

ruta_pfs_p='../Classification/2-Classification/'
ruta_pfs_f='../Classification/4-Future_Classification/'

PFS_p = xr.open_dataset(ruta_pfs_p+'PFS.nc')['PFS']
PFS_f = xr.open_dataset(ruta_pfs_f+'PFS.nc')['PFS']
PFS = xr.open_dataset(ruta_pfs_f+'dPFS.nc')['PFS']
PFS_p.values[PFS_p.values==0]=np.nan
PFS_f.values[PFS_f.values==0]=np.nan


m = Basemap(epsg=4326, resolution='i',	
				llcrnrlat=-90., urcrnrlat=90.,	
				llcrnrlon=-180., urcrnrlon=180.)
lats = PFS['lat']
lons = PFS['lon']
x, y = np.meshgrid(lons.values, lats.values)


fig = plt.figure(figsize=(6,4))
gs = fig.add_gridspec(3, 2)
ax1 = fig.add_subplot(gs[0:-1,:])
m.drawcoastlines(linewidth=0.25,zorder=3)
m.drawcountries(linewidth=0.1,zorder=3)
m.drawlsmask(land_color='none',ocean_color="#e4e4e4",lakes=True,zorder=2)
clevs=np.array([-4.,-3,-2.,-1.,1.,2.,3.,4.])
cmap = cmocean.cm.delta #mcolors.ListedColormap(["#40a4de", "#328498", "#4fb93d", "#83a10c", "#ffdd2e", "#ea7a17", "#d82e2e"])
norm = mcolors.BoundaryNorm(clevs, cmap.N, extend='both')    
cs2 = m.contourf(x, y, PFS, clevs, cmap=cmap, norm=norm, alpha=1., extend='both', zorder=1)  	
axins = inset_axes(ax1,
                   width="2%",  
                   height="85%",  
                   loc='lower left',
                   bbox_to_anchor=(1.01, 0., 1, 1),
                   bbox_transform=ax1.transAxes,
                   borderpad=0,
                   )
cbar = fig.colorbar(cs2,ticks=clevs,cax=axins) 							
ax1.text(0.01, 0.99, 'a',family='sans-serif',weight='bold',size=7, horizontalalignment='left', verticalalignment='top',transform=ax1.transAxes)
ax1.text(1.01, 0.99, '$\Delta$PFS \n$(months)$',family='sans-serif',weight='bold',size=7, horizontalalignment='left', verticalalignment='top',transform=ax1.transAxes)
plt.gcf()

ax2 = fig.add_subplot(gs[2,0])
m.drawcoastlines(linewidth=0.25,zorder=3)
m.drawcountries(linewidth=0.1,zorder=3)
m.drawlsmask(land_color='none',ocean_color="#e4e4e4",lakes=True,zorder=2)
clevs=np.array([0,1,2,3,4,5,6])
cmap = cmocean.cm.thermal #mcolors.ListedColormap(["#40a4de", "#328498", "#4fb93d", "#83a10c", "#ffdd2e", "#ea7a17", "#d82e2e"])
norm = mcolors.BoundaryNorm(clevs, cmap.N, extend='max')    
cs2 = m.contourf(x, y, PFS_p, clevs, cmap=cmap, norm=norm, alpha=1., extend='max', zorder=1)  	
axins = inset_axes(ax2,
                   width="2%",  
                   height="80%",  
                   loc='lower left',
                   bbox_to_anchor=(1.01, 0., 1, 1),
                   bbox_transform=ax2.transAxes,
                   borderpad=0,
                   )
cbar = fig.colorbar(cs2,ticks=clevs,cax=axins) 							
ax2.text(0.01, 0.99, 'b',family='sans-serif',weight='bold',size=7,horizontalalignment='left', verticalalignment='top',transform=ax2.transAxes)
ax1.text(1.01, 0.99, 'PFS \n$(months)$',family='sans-serif',weight='bold',size=5, horizontalalignment='left', verticalalignment='top',transform=ax2.transAxes)
plt.gcf()

ax3 = fig.add_subplot(gs[2,1])
m.drawcoastlines(linewidth=0.25,zorder=3)
m.drawcountries(linewidth=0.1,zorder=3)
m.drawlsmask(land_color='none',ocean_color="#e4e4e4",lakes=True,zorder=2)
clevs=np.array([0,1,2,3,4,5,6])
cmap = cmocean.cm.thermal #mcolors.ListedColormap(["#40a4de", "#328498", "#4fb93d", "#83a10c", "#ffdd2e", "#ea7a17", "#d82e2e"])
norm = mcolors.BoundaryNorm(clevs, cmap.N, extend='max')    
cs2 = m.contourf(x, y, PFS_f, clevs, cmap=cmap, norm=norm, alpha=1., extend='max',zorder=1)  	
axins = inset_axes(ax3,
                   width="2%", 
                   height="80%",  
                   loc='lower left',
                   bbox_to_anchor=(1.01, 0., 1, 1),
                   bbox_transform=ax3.transAxes,
                   borderpad=0,
                   )
cbar = fig.colorbar(cs2,ticks=clevs,cax=axins) 							
ax3.text(0.01, 0.99, 'c',family='sans-serif',weight='bold',size=7,horizontalalignment='left', verticalalignment='top',transform=ax3.transAxes)
ax1.text(1.01, 0.99, 'PFS \n$(months)$',family='sans-serif',weight='bold',size=5, horizontalalignment='left', verticalalignment='top',transform=ax3.transAxes)
plt.gcf()

plt.subplots_adjust(top = 0.95, bottom=0.05, hspace=0.1, wspace=0.25)

fig.savefig('Figure3.png',dpi=300, bbox_inches='tight') 						
#fig.savefig('Figure3.pdf', bbox_inches='tight')                         
